import sqlite3
from typing import Generator
from contextlib import contextmanager

DATABASE_URL = "D:\\Attendrix.db"

@contextmanager
def get_db() -> Generator[sqlite3.Connection, None, None]:
    """获取数据库连接的上下文管理器"""
    conn = sqlite3.connect(DATABASE_URL)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()

def get_db_connection() -> sqlite3.Connection:
    """获取数据库连接，不使用上下文管理器
    
    注意：调用者负责关闭连接
    """
    conn = sqlite3.connect(DATABASE_URL)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    """初始化数据库，创建必要的表"""
    with get_db() as conn:
        with open("考勤系统数据库结构.sql", "r", encoding="utf-8") as f:
            conn.executescript(f.read())
        conn.commit() 